package io.hepu.robotize.security;

import com.auth0.jwt.exceptions.JWTVerificationException;
import com.auth0.jwt.exceptions.TokenExpiredException;
import io.hepu.robotize.IConstant;
import io.hepu.robotize.IContext;
import io.hepu.robotize.auth.JRedis;
import io.hepu.robotize.auth.Jwts;
import io.hepu.robotize.enums.SC;
import io.hepu.robotize.exception.AuthGrantException;
import io.hepu.robotize.util.Jacksons;
import io.hepu.robotize.util.Limits;
import org.apache.shiro.web.filter.authc.BasicHttpAuthenticationFilter;
import org.apache.shiro.web.util.WebUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.RequestMethod;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * Shiro过滤器，对用户的JWT TOKEN进行基础验证
 * 执行流程：preHandle -> isAccessAllowed -> isLoginAttempt -> executeLogin
 */
public class JwtAuthcFilter extends BasicHttpAuthenticationFilter {
    private static final Logger LOG = LoggerFactory.getLogger(JwtAuthcFilter.class);

    @Override
    protected boolean isAccessAllowed(ServletRequest request, ServletResponse response, Object mappedValue) {
        LOG.debug("Authentication => Allowed: {}", IContext.getRequest().getRequestURL());
        // Access verification
        if (this.isLoginAttempt(request, response)) {
            try {
                return this.executeLogin(request, response);
            } catch (Exception e) {
                Throwable throwable = e.getCause();
                /* JWT token expired */
                if (throwable instanceof TokenExpiredException) {
                    return this.tryRefreshToken(request, response);
                }
                /* Signature resulted invalid or decode failure */
                if (throwable instanceof JWTVerificationException) {
                    this.onException(request, response, AuthGrantException.AUTHENTICATED_INVALID);
                }
                /* AuthGrantException failure */
                if (e instanceof AuthGrantException) {
                    this.onException(request, response, (AuthGrantException) e);
                }
                return false;
            }
        } else {
            // 限制所有未授权请求除了匿名访问过滤器之外：DefaultShiroFilterChainDefinition
            this.onException(request, response, AuthGrantException.UNAUTHORIZED);
            return false;
        }
    }

    private void onException(ServletRequest request, ServletResponse response, AuthGrantException ae) {
        LOG.error("AuthGrantException caught: code={}, message={}", ae.getCode(), ae.getMessage());
        HttpServletResponse resp = (HttpServletResponse) response;
        resp.setStatus(SC.FAILURE);
        resp.setCharacterEncoding(IConstant.CHARSET_UTF8);
        resp.setContentType(IConstant.APPLICATION_JSON);
        try (PrintWriter out = resp.getWriter()) {
            out.write(Jacksons.serialize(ae));
            out.flush();
        } catch (IOException ioe) {
            LOG.error("Failed to write response: {}", ioe.getMessage());
        }
    }

    @Override
    protected boolean onAccessDenied(ServletRequest request, ServletResponse response) {
        LOG.debug("Authentication => Denied: {}", IContext.getRequest().getRequestURL());
        this.sendChallenge(request, response);
        return false;
    }

    @Override
    protected boolean isLoginAttempt(ServletRequest request, ServletResponse response) {
        String authorization = this.getAuthzHeader(request);
        return authorization != null;
    }

    @Override
    protected boolean executeLogin(ServletRequest request, ServletResponse response) throws Exception {
        LOG.trace("Authentication => Scheme: {}, Platform: {}", this.getAuthzScheme(), IContext.getAuthzPlatform());
        String authzHeader = Jwts.trimScheme(this.getAuthzHeader(request));
        JwtToken token = new JwtToken(authzHeader, IContext.getAuthzPlatform());
        this.getSubject(request, response).login(token);
        return true;
    }

    public boolean tryRefreshToken(ServletRequest request, ServletResponse response) {
        String tokenExpired = Jwts.trimScheme(this.getAuthzHeader(request));
        String userId = Jwts.getSubject(tokenExpired);
        String timestampKey = IContext.getJwtRefreshKey(userId);
        if (JRedis.exists(timestampKey)) {
            String redisTimestamp = JRedis.string(timestampKey);
            Long jwtTimestamp = Jwts.getSigningAt(tokenExpired);
            if (Limits.equals(redisTimestamp, jwtTimestamp)) {
                long latest = System.currentTimeMillis();
                // Set the latest timestamp to redis as refresh token expiration time
                JRedis.serialize(timestampKey, latest, IConstant.JWT_REFRESH_TOKEN_EXPIRATION);
                // Set the latest timestamp to JWT as signing time
                String tokenRefreshed = Jwts.sign(tokenExpired, latest);
                LOG.info("::: REFRESH JwtToken = {} ", tokenRefreshed);
                JwtToken jwtToken = new JwtToken(tokenRefreshed, IContext.getAuthzPlatform());
                // Login to Shiro again with the refreshed token.
                // Exception will be caught if be thrown, if no exception is thrown, it means login is successful
                this.getSubject(request, response).login(jwtToken);
                // Store the refreshed token in the Authorization field to the Response Header and return.
                HttpServletResponse resp = WebUtils.toHttp(response);
                resp.setHeader("Authorization", tokenRefreshed);
                return true;
            }
        }
        return false;
    }

    @Override
    protected boolean preHandle(ServletRequest request, ServletResponse response) throws Exception {
        HttpServletRequest req = (HttpServletRequest) request;
        HttpServletResponse resp = (HttpServletResponse) response;
        resp.setHeader("Access-control-Allow-Origin", req.getHeader("Origin"));
        resp.setHeader("Access-Control-Allow-Methods", "GET,POST,OPTIONS,PUT,DELETE");
        resp.setHeader("Access-Control-Allow-Headers", req.getHeader("Access-Control-Request-Headers"));
        resp.setHeader("Access-Control-Allow-Credentials", "true");
        resp.setHeader("Access-Control-Max-Age", "3600");
        // 对于option请求，直接返回正常状态
        if (req.getMethod().equals(RequestMethod.OPTIONS.name())) {
            resp.setStatus(SC.SUCCESS);
            return false;
        }
        return super.preHandle(request, response);
    }
}
